package com.tfswx.jcyw.webssh.service.impl;

import com.alibaba.fastjson.JSON;
import com.jcraft.jsch.Channel;
import com.jcraft.jsch.JSch;
import com.jcraft.jsch.JSchException;
import com.tfswx.jcyw.webssh.bean.SSHConnectInfo;
import com.tfswx.jcyw.webssh.bean.WebSSHData;
import com.tfswx.jcyw.webssh.constant.ConstantPool;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;

import javax.websocket.OnClose;
import javax.websocket.OnMessage;
import javax.websocket.OnOpen;
import javax.websocket.Session;
import javax.websocket.server.ServerEndpoint;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.nio.charset.StandardCharsets;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.stream.Collectors;

/**
 * @author 张述江
 * @date 2023-11-07
 */
@Slf4j
@Component
@ServerEndpoint("/webssh")
public class WebSocketServer {

    private final Map<String, SSHConnectInfo> sshMap = new ConcurrentHashMap<>();
    /**
     * 线程池
     */
    private final ExecutorService executorService = Executors.newCachedThreadPool();

    @OnOpen
    public void onOpen(Session websocketSession) {
        String id = websocketSession.getId();
        System.out.println("建立了连接：" + id);
    }

    @OnMessage
    public void onMessage(String message, Session websocketSession) {
        System.out.println(websocketSession.getId() + "发来消息：" + message);
        WebSSHData webSSHData = JSON.parseObject(message, WebSSHData.class);
        System.out.println(webSSHData);
        if (ConstantPool.WEBSSH_OPERATE_CONNECT.equals(webSSHData.getOperate())) {
            //如果是连接请求
            connectSsh(websocketSession, webSSHData);
        } else if (ConstantPool.WEBSSH_OPERATE_COMMAND.equals(webSSHData.getOperate())) {
            //如果是发送命令的请求
            //发送命令到终端
            transToSSH(webSSHData);
        } else if (ConstantPool.WEBSSH_OPERATE_DISCONNECT.equals(webSSHData.getOperate())) {
            closeSsh(webSSHData, true);
        } else {
            log.error("不支持的操作");
            close(websocketSession, true);
        }
    }

    @OnClose
    public void onClose(Session websocketSession) {
        System.out.println("关闭连接：" + websocketSession.getId());
        close(websocketSession, false);
    }

    @SneakyThrows
    public void connectSsh(Session websocketSession, WebSSHData webSSHData) {
        log.info("开始初始化ssh连接");
        JSch jsch = new JSch();
        jsch.addIdentity("~/.ssh/id_rsa");
        SSHConnectInfo sshConnectInfo = new SSHConnectInfo();
        sshConnectInfo.setJsch(jsch);
        sshConnectInfo.setWebSocketSession(websocketSession);
        sshConnectInfo.setSshid(webSSHData.getId());
        //将这个ssh连接信息放入map中
        sshMap.put(sshConnectInfo.getSshid(), sshConnectInfo);

        //启动线程异步处理
        executorService.execute(() -> {
            try {
                //连接到终端
                connectToSSH(sshConnectInfo, webSSHData);
            } catch (JSchException | IOException e) {
                log.error("webssh连接异常", e);
                log.error("异常信息:{}", e.getMessage());
                closeSsh(webSSHData.getId(), true);
            }
        });
    }

    private void connectToSSH(SSHConnectInfo sshConnectInfo, WebSSHData webSSHData) throws JSchException, IOException {
        com.jcraft.jsch.Session session = null;
        Properties config = new Properties();
        config.put("StrictHostKeyChecking", "no");
        //获取jsch的会话
        session = sshConnectInfo.getJsch().getSession(webSSHData.getUsername(), webSSHData.getHost(), webSSHData.getPort());
        sshConnectInfo.setSshSession(session);
        session.setConfig(config);
        //设置密码
        session.setPassword(webSSHData.getPassword());

        Session webSocketSession = sshConnectInfo.getWebSocketSession();
        try {
            //连接  超时时间30s
            session.connect(30000);
        } catch (Exception e) {
            sendMessage(webSocketSession, "\r\nError 连接失败：" + e.getMessage(), webSSHData.getId());
            closeSsh(webSSHData.getId(), true);
            return;
        }

        //开启shell通道
        Channel channel = session.openChannel("shell");

        //通道连接 超时时间3s
        channel.connect(3000);

        //设置channel
        sshConnectInfo.setChannel(channel);

        //转发消息
        sendMessage(webSocketSession, "成功连接到【" + webSSHData.getHost() + "】服务器\r\n", webSSHData.getId());

        //读取终端返回的信息流
        try (InputStream inputStream = channel.getInputStream()) {
            //循环读取
            byte[] buffer = new byte[1024];
            int i;
            //如果没有数据来，线程会一直阻塞在这个地方等待数据。
            while ((i = inputStream.read(buffer)) != -1) {
                sendMessage(webSocketSession, new String(Arrays.copyOfRange(buffer, 0, i), StandardCharsets.UTF_8), webSSHData.getId());
            }

        } finally {
            //断开连接后关闭会话
            closeSsh(webSSHData.getId(), true);
        }
    }

    public void sendMessage(Session webSocketSession, String str, String id) throws IOException {
        Map<String, String> map = new HashMap<>();
        map.put("id", id);
        map.put("data", str);
        String s = JSON.toJSONString(map);
        log.info("开始发送消息到客户端");
        System.out.println(map);
        synchronized (webSocketSession.getId()) {
            if (webSocketSession.isOpen()) {
                try {

                    webSocketSession.getBasicRemote().sendText(s);
                } catch (Exception e) {
                    e.printStackTrace();
                }
            }
        }
        log.info("完成发送消息到客户端");
    }

    public void close(Session session, boolean sendMsg) {
        //获取随机生成的uuid
        String socketId = session.getId();
        List<SSHConnectInfo> sshList = sshMap.values().stream().filter(x -> socketId.equals(x.getSocketId())).collect(Collectors.toList());

        for (SSHConnectInfo sshConnectInfo : sshList) {
            closeSsh(sshConnectInfo.getSshid(), sendMsg);
        }

        try {
            session.close();
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    private void closeSsh(WebSSHData webSSHData, boolean sendMsg) {
        String sshid = webSSHData.getId();
        closeSsh(sshid, sendMsg);
    }

    @SneakyThrows
    private void closeSsh(String sshid, boolean sendMsg) {
        SSHConnectInfo sshConnectInfo = sshMap.get(sshid);
        if (sshConnectInfo == null) {
            return;
        }
        //断开连接
        com.jcraft.jsch.Session sshSession = sshConnectInfo.getSshSession();
        if (sshSession != null && sshSession.isConnected()) {
            log.info("正在断开连接：{},{}", sshid, sshSession.getHost());
            sshSession.disconnect();
            if (sendMsg) {
                Session webSocketSession = sshConnectInfo.getWebSocketSession();
                sendMessage(webSocketSession, "\r\n连接已关闭", sshid);
            }
        }

        //map中移除该ssh连接信息
        sshMap.remove(sshid);

    }

    private void transToSSH(WebSSHData webSSHData) {

        String command = webSSHData.getCommand();
        SSHConnectInfo sshConnectInfo = sshMap.get(webSSHData.getId());
        if (sshConnectInfo != null) {
            Session session = sshConnectInfo.getWebSocketSession();
            try {
                Channel channel = sshConnectInfo.getChannel();
                //发送命令到终端
                if (channel != null) {
                    if (channel.isConnected()) {
                        OutputStream outputStream = channel.getOutputStream();
                        outputStream.write(command.getBytes());
                        outputStream.flush();
                    } else if (channel.isClosed()) {
                        log.warn("通道已关闭");
                    }
                }
            } catch (IOException e) {
                log.error("webssh连接异常", e);
                log.error("异常信息:{}", e.getMessage());
                closeSsh(webSSHData.getId(), true);
            }
        }
    }
}
